#!/usr/bin/python
# -*- coding: utf8 -*-

from pyinmune import modelo, agente, estado, poblacion, plots
from datosPneumo import porcentajeEnfermos, cantDePersonas

from collections import defaultdict
import random

import psyco
psyco.full()

EstadoMuerto = agente.EstadoMuerto
EstadoSano = estado.Estado('Sano')
EstadoEnfermo = estado.Estado('Enfermo')
EstadoInmune = estado.Estado('Inmune')

CocienteDePoblacion = 100


class AgentePneumonia(agente.Agente):
            
    def inicializar(self):
        
        self.cocientesContagio = porcentajeEnfermos
        # Peso esperado de los ejes (para distribuir las probabilidades de contagio)
        self.pesoDeEjes = 120.0
        # Diccionario de vacuna a número de iteracion donde se lo vacunó
        self.fechasDeVacunaciones = {}
        
    def encontrarseCon(self, otro, peso):
        if self.estado == EstadoSano:
            if otro.estado == EstadoEnfermo:
                # a es el threshold para enfermarse
                a = 50*(peso * self.cocientesContagio[self['edad']][self['sexo']] / self.pesoDeEjes)
                
                if self.vacunado:
                    for each in self.fechasDeVacunaciones:
                        a -=0.009
                    a = max(a,0)    
                
                if random.random() < a:
                    self.cambiarEstado(EstadoEnfermo)
                    
    def terminarIteracion(self):
        # Sacamos el efecto de las vacunas
        for each in self.fechasDeVacunaciones.keys():
            self.fechasDeVacunaciones[each] -= 1
            if self.fechasDeVacunaciones[each] == 0:
                del self.fechasDeVacunaciones[each]
                
        if self.fechasDeVacunaciones == {}:
            self.vacunado=False
            
        if self.estado == EstadoEnfermo:
            if self.iteracionesEnEstado > 7:
                self.cambiarEstado(EstadoMuerto)
                
            if random.random()*self.iteracionesEnEstado >= 1.95:
                self.cambiarEstado(EstadoInmune)
                
        if self.estado == EstadoInmune \
            and self.iteracionesEnEstado > 15:
            self.cambiarEstado(EstadoSano)
            

class AislamientoPreventivo(modelo.Aislamiento):
    
    def criterioDeAislamientoPoblacion(self,poblacion):
        return len(EstadoEnfermo) > len(EstadoSano) * 0.1 or \
              (len(EstadoEnfermo) < 20 and poblacion.iteracion > 10)
    
    def criterioDeDesaislamientoPoblacion(self,poblacion):
        return True
    
    def criterioDeAislamientoAgente(self,agente,vecinos):
        return random.random() < 0.1 and \
                agente.estado == EstadoEnfermo and \
                len([x for x in vecinos if x.estado == EstadoSano]) > 0.7*len(vecinos) and \
                not agente.aislado
        
    
    def criterioDeDesaislamientoAgente(self,agente,vecinos):
        return agente.estado == EstadoInmune or agente.estado == EstadoMuerto


class Vacuna1(modelo.Vacuna):
    
    def inicializar(self):
        self.duracion = 20
    
    def postVacunacion(self,agente,vecinos):
        agente.fechasDeVacunaciones[self] = self.duracion
    
    def criterioDeVacunacionPoblacion(self,poblacion):
        return poblacion.iteracion > 10 and poblacion.iteracion < 100
                  
    def criterioDeVacunacionAgente(self,agente,vecinos):
        return agente.estado == EstadoSano and \
                random.random() < 0.05 and \
                not self in agente.vacunas


class RifleSanitario(modelo.Curbing):
    
    def criterioDeCurbingPoblacion(self,poblacion):
        return poblacion.iteracion > 20 and len(EstadoEnfermo) < 100
    
    def criterioDeCurbingAgente(self,agente,vecinos):
        return agente.estado == EstadoEnfermo


class GeneradorDePoblacion(object):
    
    def __init__(self):
        self.poblacion = cantDePersonas

    def generar(self):
        Eje = poblacion.Eje
        Poblacion = poblacion.Poblacion
        print "Generando agentes...  ",
        agentes = list()
        porcentaje = 225.0/151193
        
        for rango in self.poblacion:
            for sexo in self.poblacion[rango]:
                for elem in xrange(self.poblacion[rango][sexo]/CocienteDePoblacion):
                    if random.random() <= porcentaje: 
                        agentes.append(AgentePneumonia(EstadoEnfermo, {'edad':rango,'sexo':sexo}))
                    else:
                        agentes.append(AgentePneumonia(EstadoSano, {'edad':rango,'sexo':sexo}))
        print "OK!"
        print "Generando vínculos... ",
        vinculos = set()
        con_familia =list()
        cant = random.choice([2,3,4,5,6,7])
        
        while cant <= len(agentes):
            fam = []
            while cant > 0:
                e = random.choice(agentes)
                fam.append(e)
                agentes.remove(e)
                cant-=1
                con_familia.append(e)
                
            for each in fam:
                for each2 in fam:
                    if not (each == each2):
                        vinculos.add(Eje(each,each2,10))
            cant = random.choice([2,3,4,5,6,7])
        
        else:
            cant = len(agentes)
            fam = []
            
            while cant > 0:
                e = random.choice(agentes)
                fam.append(e)
                agentes.remove(e)
                cant-=1
                con_familia.append(e)
                
            for each in fam:
                for each2 in fam:
                    if not (each == each2):
                        vinculos.add(Eje(each,each2,10))
        

        for each in con_familia:
            cant = int(random.normalvariate(10,2))
            while cant > 0:
                tipito = random.choice(con_familia)
                if tipito != each:
                    vinculos.add(Eje(each,tipito,5))
                cant -= 1
        
        for each in con_familia:
            cant = int(random.normalvariate(20,4))
            while cant > 0:
                tipito = random.choice(con_familia)
                if tipito != each:
                    vinculos.add(Eje(each,tipito,1))
                cant -= 1
                
        print "OK!"

        return poblacion.Poblacion(con_familia, vinculos)

def main(iteraciones=100):
    # Para la barra de progreso
    from enthought.pyface.progress_dialog import ProgressDialog
    d = ProgressDialog(title="Simulación en curso", max=iteraciones, show_time=True)
    d.open()

    poblacion = GeneradorDePoblacion().generar()
    m = modelo.Modelo(poblacion)

    v = Vacuna1('tito')
    a = AislamientoPreventivo()
    r = RifleSanitario()

    m.registrarCallbacks(v.aplicarVacuna)
    m.registrarCallbacks(a.aplicarAislamiento)
    m.registrarCallbacks(a.aplicarDesaislamiento)
    m.registrarCallbacks(r.AplicarCurbing)
    
    datos = defaultdict(lambda: [])

    datos['Sanos'].append(len(EstadoSano))
    datos['Enfermos'].append(len(EstadoEnfermo))
    datos['Muertos'].append(len(EstadoMuerto))
    datos['Inmunes'].append(len(EstadoInmune))
    
    for i in xrange(iteraciones):
        m.avanzarIteracion()
        print "Iteración:", i
        datos['Sanos'].append(len(EstadoSano))
        datos['Enfermos'].append(len(EstadoEnfermo))
        datos['Muertos'].append(len(EstadoMuerto))
        datos['Inmunes'].append(len(EstadoInmune))

        d.update(i)

        if len(EstadoEnfermo) == 0:
            print "No hay más enfermos - terminando simulación."
            break

    iteraciones_reales = i+1

    d.update(iteraciones)
    d.close()
    
    datos['Vacunados'] = [sum((len(v.vacunados[j]) for j in range(i+1))) for i in range(iteraciones_reales+1)]
    datos['Aislados'] = [sum((len(a.aislados[j]) for j in range(i+1))) for i in range(iteraciones_reales+1)]
    print datos

    plots.plotear_enfermedad(datos)

if __name__ == '__main__':
    main()
